import cv2
import glob
import os
import numpy as np
import argparse


def parse_args():
    parser = argparse.ArgumentParser(description='use sobel to get the edge of images')

    parser.add_argument('--input', default='work_dirs/Train_resized')
    parser.add_argument('--output', default='work_dirs/Train_resized_edge')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    if not os.path.exists(args.output):
        os.mkdir(args.output)

    for img_dir in glob.glob(os.path.join(args.input, '*.jpg')):
        img = cv2.imread(img_dir)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 二值化
        _, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

        # sobel 水平方向边缘检测
        edges = cv2.Sobel(img, cv2.CV_16S, 1, 0)
        edgesh = cv2.convertScaleAbs(edges)

        # sobel 竖直方向边缘检测
        edges = cv2.Sobel(img, cv2.CV_16S, 0, 1)
        edgesv = cv2.convertScaleAbs(edges)

        sobelxy = cv2.addWeighted(edgesh, 0.5, edgesv, 0.5, 0)
        cv2.imwrite(os.path.join(args.output, 'sobel-'+img_dir.split('/')[-1]), sobelxy)

if __name__ == '__main__':
    main()
